package org.finesys.chat.core.base.embedding.factory;

import lombok.extern.slf4j.Slf4j;
import org.finesys.ai.constants.EmbedStoreEnum;
import org.finesys.ai.entity.AigcEmbedStore;
import org.finesys.chat.core.base.embedding.EmbeddingStore;
import org.finesys.chat.core.base.embedding.store.milvus.MilvusEmbeddingStore;
import org.finesys.chat.core.base.embedding.store.pgvector.PgVectorEmbeddingStore;
import org.finesys.chat.core.base.embedding.store.redis.RedisEmbeddingStore;
import org.finesys.chat.core.base.segment.TextSegment;
import org.springframework.stereotype.Component;
import org.springframework.util.StringUtils;

import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

@Slf4j
@Component
public class EmbeddingStoreFactory {

    private final Map<String, EmbeddingStore<TextSegment>> embeddingStores = new ConcurrentHashMap<>();


    /**
     * 初始话方法
     */
    public void init(List<AigcEmbedStore> embedStoreList) {
        // 停止之前的嵌入存储服务
        if (!embeddingStores.isEmpty()) {
            embeddingStores.values().stream().forEach(EmbeddingStore::stop);
        }
        //重新初始化嵌入存储服务
        embedStoreList.forEach(this::addOrUpdate);
    }

    /**
     * 添加或更新嵌入存储服务
     *
     * @param embedStore 嵌入存储配置对象
     */
    public void addOrUpdate(AigcEmbedStore embedStore) {
        // 停止之前的嵌入存储服务
        if (embeddingStores.containsKey(embedStore.getId())) {
            embeddingStores.get(embedStore.getId()).stop();
        }
        //重新初始化嵌入存储服务
        String provider = embedStore.getProvider();
        //redis
        if (EmbedStoreEnum.REDIS.name().equals(provider)) {
            RedisEmbeddingStore.Builder builder = RedisEmbeddingStore.builder();
            builder.host(embedStore.getHost())
                    .port(embedStore.getPort())
                    .indexName(embedStore.getDatabaseName())
                    .dimension(embedStore.getDimension());
            if (StringUtils.hasText(embedStore.getUserName()) && StringUtils.hasText(embedStore.getPassword())) {
                builder.user(embedStore.getUserName()).password(embedStore.getPassword());
            }
            RedisEmbeddingStore redisEmbeddingStore = builder.build();
            embeddingStores.put(embedStore.getId(), redisEmbeddingStore);
            return;
        }
        //milvus
        if (EmbedStoreEnum.MILVUS.name().equals(provider)) {
            MilvusEmbeddingStore.Builder builder = MilvusEmbeddingStore.builder();
            builder.host(embedStore.getHost()).port(embedStore.getPort())
                    .username(embedStore.getUserName()).password(embedStore.getPassword())
                    .databaseName(embedStore.getDatabaseName())
                    .dimension(embedStore.getDimension())
                    .collectionName(embedStore.getTableName());
            MilvusEmbeddingStore milvusEmbeddingStore = builder.build();
            embeddingStores.put(embedStore.getId(), milvusEmbeddingStore);
            return;
        }
        //Pgvector
        if (EmbedStoreEnum.PGVECTOR.name().equals(provider)) {
            PgVectorEmbeddingStore.Builder builder = PgVectorEmbeddingStore.builder();
            builder.host(embedStore.getHost()).port(embedStore.getPort())
                    .database(embedStore.getDatabaseName())
                    .table(embedStore.getTableName())
                    .dimension(embedStore.getDimension())
                    .user(embedStore.getUserName())
                    .password(embedStore.getPassword())
                    .indexListSize(1)
                    .useIndex(true)
                    .createTable(true)
                    .dropTableFirst(false);
            PgVectorEmbeddingStore pgVectorEmbeddingStore = builder.build();
            embeddingStores.put(embedStore.getId(), pgVectorEmbeddingStore);
            return;
        }
        //其他嵌入存储服务
        if (log.isErrorEnabled()) {
            log.error("向量数据库初始化失败：[{}] --- [{}]，数据库配置信息：[{}]", embedStore.getName(), embedStore.getProvider(), embedStore);
        }
    }


    public void unregister(AigcEmbedStore embedStore) {
        String id = embedStore.getId();
        if (embeddingStores.containsKey(id)) {
            EmbeddingStore<TextSegment> embeddingStore = embeddingStores.get(id);
            embeddingStore.stop();
            embeddingStores.remove(id);
        }
    }

    /**
     * 根据给定的id获取对应的嵌入存储。
     *
     * @param id 嵌入存储的唯一标识符
     * @return 与给定id对应的嵌入存储对象，如果id不存在，则返回null
     */
    public EmbeddingStore<TextSegment> get(String id) {
        return embeddingStores.get(id);
    }

    public boolean contains(String id) {
        return embeddingStores.containsKey(id);
    }
}
